﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Net;

namespace EkspSys.Network
{
    class ProcessHelper
    {
        private const int NO_ERROR = 0;

        public enum State
        {
            CLOSED, LISTEN, SENT, RECIEVED, ESTABLISHED, FIN_WAIT_1
        };

        // TODO: Set state as State type.
        public struct TcpRow
        {
            public int state;
            public IPAddress localAddress;
            public UInt16 localPort;
            public IPAddress remoteAddr;
            public UInt16 remotePort;
            public Process owningProcess;
            public int offloadState;
        }

        [DllImport("iphlpapi.dll", SetLastError = true)]
        private static extern int GetTcpTable2(byte[] table, out int size, bool order);

        private static uint readDWORD(byte[] buf, ref int pos)
        {
            pos += 4;
            return (uint)(((uint)buf[pos - 1]) << 24) + (((uint)buf[pos - 2]) << 16) + (((uint)buf[pos - 3]) << 8) + ((uint)buf[pos - 4]);
        }

        public static List<TcpRow> GetTcpConnections()
        {
            int bufSize = 20000;
            byte[] buf = new byte[bufSize];

            int res = GetTcpTable2(buf, out bufSize, true);
            if (res != NO_ERROR)
            {
                Console.WriteLine("Failed loading TCP table. It properly is because the table buffer is too small.");
            }

            int pos = 0;
            uint entries = readDWORD(buf, ref pos);
            List<TcpRow> rows = new List<TcpRow>(Convert.ToInt32(entries));
            for (int i = 0; i < entries; i++)
            {
                TcpRow row = new TcpRow();
                row.state = Convert.ToInt32(readDWORD(buf, ref pos));
                row.localAddress = new IPAddress(readDWORD(buf, ref pos));
                row.localPort = Convert.ToUInt16(readDWORD(buf, ref pos));
                row.remoteAddr = new IPAddress(readDWORD(buf, ref pos));
                row.remotePort = Convert.ToUInt16(readDWORD(buf, ref pos));
                try
                {
                    row.owningProcess = Process.GetProcessById(Convert.ToInt32(readDWORD(buf, ref pos)));
                }
                catch (ArgumentException ex)
                {
                    row.owningProcess = null;
                }
                pos += 4;

                rows.Add(row);
            }

            return rows;
        }

        public static Process GetProcessForPort(Protocol proto, int lcalPrt)
        {
            if (proto == Protocol.TCP)
            {
                List<TcpRow> rows = GetTcpConnections();
                foreach (TcpRow row in rows)
                {
                    if (row.localPort == lcalPrt)
                        return row.owningProcess;
                }
                return null;
            }
            else if (proto == Protocol.UDP)
            {
                // TODO: Add support for UDP
                return null;
            }
            else
            {
                throw new ApplicationException("Unsupported protocol!");
            }
        }
    }
}
